package comCrypto

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
)

// GenerateRandomKey 获取随机数
func GenerateRandomKey(keySize int) ([]byte, error) {
	if keySize != 16 && keySize != 24 && keySize != 32 {
		return nil, errors.New("invalid key size (must be 16/24/32)")
	}

	key := make([]byte, keySize)
	if _, err := rand.Read(key); err != nil {
		return nil, fmt.Errorf("failed to generate key: %w", err)
	}
	return key, nil
}

// AESGCMEncrypt 使用 AES-GCM 加密数据
func AESGCMEncrypt(plaintext []byte, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher block: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}

	// 组合 nonce + ciphertext
	return gcm.Seal(nonce, nonce, plaintext, nil), nil
}

// AESGCMDecrypt 使用 AES-GCM 解密数据
func AESGCMDecrypt(ciphertext []byte, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("failed to create cipher block: %w", err)
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	nonceSize := gcm.NonceSize()
	if len(ciphertext) < nonceSize {
		return nil, errors.New("invalid ciphertext length")
	}

	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]
	return gcm.Open(nil, nonce, ciphertext, nil)
}

func main() {
	// 示例用法
	message := []byte("这是一段需要加密的敏感数据")

	// 1. 生成密钥（推荐256位）
	key, err := GenerateRandomKey(32)
	if err != nil {
		panic(err)
	}

	// 2. 加密数据
	ciphertext, err := AESGCMEncrypt(message, key)
	if err != nil {
		panic(err)
	}

	// 3. 输出加密结果（Base64编码便于存储）
	fmt.Printf("加密结果: %s\n", base64.StdEncoding.EncodeToString(ciphertext))

	// 4. 解密数据
	decrypted, err := AESGCMDecrypt(ciphertext, key)
	if err != nil {
		panic(err)
	}

	fmt.Printf("解密结果: %s\n", decrypted)
}
